import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import find_peaks


# 定义MUSIC算法函数
def music_algorithm(signal_array, sensor_distance, wavelength, theta_resolution, num_signals):
    # 获取信号个数和传感器数量
    num_sensors, num_samples = signal_array.shape
    # 计算空间参数
    omega = 2 * np.pi * sensor_distance / wavelength * np.cos(np.deg2rad(np.arange(0, 180, theta_resolution)))
    # 计算信号协方差矩阵
    covariance_matrix = np.cov(signal_array)
    # 对协方差矩阵进行特征值分解
    eigenvalues, eigenvectors = np.linalg.eigh(covariance_matrix)  # 特征值从小到大排列

    # 选取前K个特征值和对应的特征向量
    dim = num_sensors - num_signals
    noisy_subspace = eigenvectors[:, :dim]

    # 初始化谱函数
    spectrum = np.zeros(len(omega))  # len = 180

    # 计算谱函数
    for i in range(len(omega)):  # 遍历每一个谱角度
        a_theta = np.exp(1j * omega[i] * np.arange(num_sensors)).reshape(4,1)
        A = a_theta.conj().T
        B = noisy_subspace
        C = noisy_subspace.conj().T
        D = a_theta
        spectrum[i] = 1 / np.abs(np.dot(np.dot(np.dot(A, B), C), D))

    # 寻找波达方向的估计值
    # peak_index = np.argsort(spectrum)[-num_signals:]
    peak_index, _ = find_peaks(spectrum,distance=3, prominence=10)
    # estimated_theta = np.rad2deg(np.arccos(omega[peak_index] * wavelength / (2 * np.pi * sensor_distance)))

    return spectrum, peak_index

# 生成模拟水听器阵列信号
def generate_signal(num_signals, num_sensors, sensor_distance, wavelength, SNR_dB, num_samples):
    '''
    return:
    array_signal:(num_sensors, num_samples) 输出仿真信号
    '''
    noise_power = 10 ** (-SNR_dB / 10)
    signal_angles = np.random.uniform(0, np.pi, num_signals)
    for i, angle in enumerate(signal_angles):
        print(f"信号{i}的角度真值为：{np.rad2deg(angle)}")

    # 生成模拟水听器阵列信号
    array_signal = np.zeros((num_sensors, num_samples), dtype=complex)
    d_over_lambda = sensor_distance / wavelength

    for signal_idx, angle_truth in enumerate(signal_angles):
        for i in range(num_samples):
            source_signal = np.exp(1j * np.random.uniform(0, 2 * np.pi))  # 随机相位的声源信号
            noise = np.random.normal(0, np.sqrt(noise_power / 2), size=(num_sensors,)) + \
                    1j * np.random.normal(0, np.sqrt(noise_power / 2), size=(num_sensors,))
            array_signal[:, i] += np.exp(
                1j * 2 * np.pi * np.arange(num_sensors) * np.cos(angle_truth) * d_over_lambda) * source_signal + noise

    return array_signal


def draw_spectrum_result(spectrum, estimated_theta, theta_resolution=1):
    # 绘制声源定向结果
    plt.figure()
    plt.plot(np.arange(0, 180, theta_resolution), spectrum)
    for i, angle in enumerate(estimated_theta):
        print(f"估计的角度{i}为：{angle}")
        plt.axvline(x=angle, color='r', linestyle='--', linewidth=1)
        plt.text(angle+12, -np.max(spectrum)/25, f'x={angle}', color='r', fontsize=12, ha='center')
        plt.plot(angle, spectrum[angle], "o", label="peaks")
    plt.title('MUSIC Algorithm for Source Localization')
    plt.xlabel('Theta (degrees)')
    plt.ylabel('Spectrum')
    plt.grid(True)
    plt.show()